/*
 * Copyright (C) 2012 Roman Yakovenko <b179a9@gmail.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
 * of the Software, and to permit persons to whom the Software is furnished to do
 * so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

package com.b17.crypto.provider;


import com.b17.crypto.core.Utils;

/**
 * Created by b17
 * 1/31/12 11:20 PM
 */
public class GOST implements Cipher {

    private static final int BLOCK_LENGTH = 8;
    public static final int KEY_LENGTH = 32;
    private static final int ROUNDS = 32;
    public static final int ENCRYPT_MODE = 0;
    public static final int DECRYPT_MODE = 1;
    private static final int ROUND_KEY_LENGTH = 4;
    private static final int SUB_BLOCK_LENGTH = 4;

    private static final byte[][] S_BOX = new byte[][]
            {
                    {4, 10, 9, 2, 13, 8, 0, 14, 6, 11, 1, 12, 7, 15, 5, 3},
                    {14, 11, 4, 12, 6, 13, 15, 10, 2, 3, 8, 1, 0, 7, 5, 9},
                    {5, 8, 1, 13, 10, 3, 4, 2, 14, 15, 12, 7, 6, 0, 9, 11},
                    {7, 13, 10, 1, 0, 8, 9, 15, 14, 4, 6, 12, 11, 2, 5, 3},
                    {6, 12, 7, 1, 5, 15, 13, 8, 4, 10, 9, 14, 0, 3, 11, 2},
                    {4, 11, 10, 0, 7, 2, 1, 13, 3, 6, 8, 5, 9, 12, 15, 14},
                    {13, 11, 4, 1, 3, 15, 5, 9, 0, 10, 14, 7, 6, 8, 2, 12},
                    {1, 15, 13, 0, 5, 7, 10, 4, 9, 2, 3, 14, 6, 11, 8, 12},
            };


    private int[] key = new int[8];


    public GOST(byte[] key) {
        init(key);
    }

    private void init(byte[] key) {
        // 256-битный ключ
        byte[] realKey = new byte[KEY_LENGTH];
        System.arraycopy(key, 0, realKey, 0, key.length);

        for (int i = 0, j = 0; i < KEY_LENGTH; i += 4, j++) {
            byte[] roundKey = new byte[ROUND_KEY_LENGTH];
            System.arraycopy(realKey, i, roundKey, 0, ROUND_KEY_LENGTH);
            this.key[j] = Utils.toInteger(roundKey);
        }
    }

    /**
     * Возвращает 64-розрядный блок данных.
     *
     * @param block
     * @param offset
     * @return
     */
    private byte[] getBlock(byte[] block, int offset) {
        int end = block.length - offset;
        byte[] result = new byte[BLOCK_LENGTH];
        if (end >= BLOCK_LENGTH) {
            System.arraycopy(block, offset, result, 0, BLOCK_LENGTH);
        } else {
            System.arraycopy(block, offset, result, 0, end);
        }
        return result;
    }

    @Override
    public void setKey(byte[] key) {
        init(key);
    }

    @Override
    public byte[] encrypt(byte[] data) {
        byte[] result = new byte[getTransformResultSize(data.length)];
        for (int i = 0; i < data.length; i += BLOCK_LENGTH) {
            byte[] block = getBlock(data, i);
            byte[] transformed = transformBlock(block, ENCRYPT_MODE);
            System.arraycopy(transformed, 0, result, i, transformed.length);
        }
        return result;
    }

    @Override
    public byte[] decrypt(byte[] data) {
        byte[] result = new byte[getTransformResultSize(data.length)];
        for (int i = 0; i < data.length; i += BLOCK_LENGTH) {
            byte[] block = getBlock(data, i);
            byte[] transformed = transformBlock(block, DECRYPT_MODE);
            System.arraycopy(transformed, 0, result, i, transformed.length);
        }
        return result;
    }

    private int getTransformResultSize(int length) {
        if (length % BLOCK_LENGTH == 0) {
            return length;
        } else {
            return (length / BLOCK_LENGTH + 1) * BLOCK_LENGTH;
        }
    }

    private byte[] transformBlock(byte[] block, int mode) {
        assert (block.length == BLOCK_LENGTH);

        byte[] leftPart = new byte[SUB_BLOCK_LENGTH];
        System.arraycopy(block, 0, leftPart, 0, SUB_BLOCK_LENGTH);
        byte[] rightPart = new byte[SUB_BLOCK_LENGTH];
        System.arraycopy(block, SUB_BLOCK_LENGTH, rightPart, 0, SUB_BLOCK_LENGTH);

        int left = Utils.toInteger(leftPart);
        int right = Utils.toInteger(rightPart);
        int tmp;
        for (int round = 0; round < ROUNDS; round++) {
            int expandedBlock = expandBlock(right, round, mode);
            left = left ^ expandedBlock;

            tmp = left;
            left = right;
            right = tmp;
        }

        tmp = left;
        left = right;
        right = tmp;

        leftPart = Utils.toBytes(left);
        rightPart = Utils.toBytes(right);

        byte[] result = new byte[BLOCK_LENGTH];
        System.arraycopy(leftPart, 0, result, 0, SUB_BLOCK_LENGTH);
        System.arraycopy(rightPart, 0, result, SUB_BLOCK_LENGTH, SUB_BLOCK_LENGTH);
        return result;
    }

    private int expandBlock(int block, int round, int mode) {
        int expandedBlock = block + getRoundKey(round, mode);
        for (int i = 0; i < 8; i++) {
            int word = (block >> (4 * i)) & 0xF;
            word = S_BOX[i][word] << (4 * i);
            expandedBlock ^= word;
        }
        expandedBlock = Utils.rol(expandedBlock, 11);
        return expandedBlock;
    }

    private int getRoundKey(int round, int mode) {
        int keyNumber = mode == ENCRYPT_MODE ? round % key.length : key.length - (round % key.length) - 1;
        if (isReversedOrder(round, mode)) {
            keyNumber = 7 - keyNumber;
        }
        return key[keyNumber];
    }

    private boolean isReversedOrder(int round, int mode) {
        if (mode == ENCRYPT_MODE) {
            return round >= 24;
        } else {
            return round >= 0 && round < 8;
        }
    }
}
